package com.demo.component.ollama.controller;

import com.demo.component.base.BaseController;
import com.demo.component.model.Question;
import com.demo.component.ollama.generate.inter.GenerateASync;
import com.demo.component.ollama.generate.inter.GenerateSync;
import com.demo.component.utils.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@Slf4j
@RestController
@RequestMapping("/ollama")
public class OllamaController extends BaseController {

    private GenerateSync generateSync;

    private GenerateASync generateASync;

    @Autowired
    public OllamaController(GenerateSync generateSync, GenerateASync generateASync) {
        this.generateSync = generateSync;
        this.generateASync = generateASync;
    }

    @PostMapping("/generate-sync")
    public String generateSync(@RequestBody Question question) {

        log.info(JsonUtil.toJson(question));

        String answer = this.generateSync.generateAnswer(question.getPrompt());

        log.info("============================");
        log.info(answer);

        return answer;
    }

    @PostMapping("/generate-async")
    public String generateASync(@RequestBody Question question) {

        log.info(JsonUtil.toJson(question));

        String answer = this.generateASync.generateAnswer(question.getPrompt());

        log.info("============================");
        log.info(answer);

        return answer;
    }

    @PostMapping("/generate-sync-image")
    public String generateSyncImage(@RequestBody Question question) {

        log.info(JsonUtil.toJson(question));

        String answer = this.generateSync.generateFromImage("c:/temp/parking.jpg", question.getPrompt());

        log.info("============================");
        log.info(answer);

        return answer;
    }

    @PostMapping("/generate-embeddings")
    public Object generateEmbeddings(@RequestBody Question question) {

        log.info(JsonUtil.toJson(question));

        List<List<Double>> embeddings = this.generateSync.generateEmbeddings(question.getPrompt());

        log.info("============================");
        System.out.println(embeddings);

        return embeddings;
    }
}
